from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded
from fastapi import Request, Depends
from auth import get_current_active_user
from typing import Dict, Any, Optional

# 创建限流器实例，基于IP地址
limiter = Limiter(key_func=get_remote_address)

# 为不同用户角色设置不同的限流策略
class RoleBasedLimiter:
    def __init__(self):
        # 定义不同角色的限流策略 (格式: 角色名 -> 限流规则)
        self.role_limits = {
            "admin": "100/minute",  # 管理员每分钟最多100次请求
            "user": "30/minute",    # 普通用户每分钟最多30次请求
            "default": "10/minute"  # 未认证用户每分钟最多10次请求
        }

    def get_rate_limit(self, user: Optional[Dict[str, Any]] = None) -> str:
        """
        根据用户角色获取对应的限流规则
        :param user: 当前用户信息
        :return: 限流规则字符串
        """
        if not user:  # 未认证用户
            return self.role_limits["default"]

        # 检查用户角色
        roles = user.get("roles", [])
        for role in roles:
            role_name = role.get("name", "")
            if role_name in self.role_limits:
                return self.role_limits[role_name]

        # 默认为普通用户权限
        return self.role_limits["user"]

# 创建基于角色的限流器实例
role_limiter = RoleBasedLimiter()

# 定义一个依赖项，用于获取当前用户的限流规则
async def get_rate_limit_for_user(user: Optional[Dict[str, Any]] = Depends(lambda: None)):    # 临时使用lambda: None替代，实际应该使用正确的依赖项
    # 注意：这里应该使用一个能正确处理用户认证状态的依赖项
    # 例如: user: Optional[Dict[str, Any]] = Depends(get_current_user_or_none)
    # 获取用户的限流规则
    rate_limit = role_limiter.get_rate_limit(user)
    # 返回限流规则
    return rate_limit

# 修复后的限流装饰器，正确传递所有参数
import functools
def rate_limit_for_user(func):
    @functools.wraps(func)
    async def wrapper(request: Request, *args, current_user: Dict[str, Any] = Depends(get_current_active_user), **kwargs):
        return await func(request, *args, current_user=current_user, **kwargs)
    return wrapper

# 定义全局的限流异常处理器
def setup_rate_limiting(app):
    app.state.limiter = limiter
    app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)
    return app